from yacs.config import CfgNode

_cfg = CfgNode()
_cfg.num_stages = 3
_cfg.debug = False
_cfg.start_stage = 1
_cfg.output_dir = "outputs/veri776_stage_pvasff"
_cfg.half = True

_cfg.train = CfgNode()
_cfg.train.epochs = [120, 120, 150]
#_cfg.data.train_size = (192, 192)
#_cfg.data.valid_size = (192, 192)
_cfg.train.modules = [['backbone', 'pam'], 
                      ['ptn', 'pam'], 
                      ['backbone', 'ptn', 'pam']]

_cfg.data = CfgNode()
_cfg.data.batch_size = [32, 8, 8]

_cfg.model = CfgNode()
_cfg.model.ckpt_period = [10, 3, 3]

_cfg.loss = CfgNode()

_cfg.optim = CfgNode()
_cfg.optim.base_lr = [0.0005, 0.0001, 0.00002]
_cfg.optim.bias_lr_factor = [1.]
_cfg.optim.base_weight_decay = [0.00005, 0.00001, 0.000005]

_cfg.scheduler = CfgNode()
_cfg.scheduler.decay_method = ["MultiStep"]
_cfg.scheduler.decay_kwargs = [["milestones", [40, 80], "gamma", 0.1]]
_cfg.scheduler.warmup_factor = [0.]
_cfg.scheduler.warmup_steps = [10]
_cfg.scheduler.warmup_method = ['linear']

_cfg.logging = CfgNode()
_cfg.logging.period = [40, 160, 160]

_cfg.test = CfgNode()
_cfg.test.period = [10, 3, 3]
_cfg.test.model_path = ["outputs/veri776_stage_pvasff/stage3/model_150.pth"]


def get_stage_configs():
    return _cfg.clone()

def get_stage_config_nodes():
    return [k for k, v in _cfg.items() if isinstance(v, CfgNode)]

def get_state_config_roots():
    return [k for k, v in _cfg.items() if not isinstance(v, CfgNode)]